#导入包

from tensorflow.keras.preprocessing.image import ImageDataGenerator

"""
 图片生成器 ：tf.keras.proprocessing.image.ImageDataGenerator
 参数：
 -rescale 输入一个整数，通常1/255,由于像素都是在0~255，rescale可以让所有像素统一乘上一个数值，如1/255，像素会被转化为0~1之间的数据
 --directory:
 --target_size:
 --batch_size
 --class_mode
"""

# def train_val_generator(data_dir, target_size, batch_size, class_mode=None, subset='training'):
#     train_val_datagen = ImageDataGenerator(rescale=1./255., validation_split=0.2)


def train_val_generator(data_dir, target_size, batch_size, class_mode=None, subset="training"):
    train_val_gen = ImageDataGenerator(rescale=1./255., validation_split=0.2)
    return train_val_gen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode,
        subset=subset
    )


def test_generator(data_dir, target_size, batch_size, class_mode=None):
    test_datagen = ImageDataGenerator(rescale=1./255.)
    return test_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode
    )


